import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import matplotlib.transforms as transforms

from load_data import dataset


def confidence_ellipse(x, y, ax, n_std=3.0,
                       facecolor='none', **kwargs):
    """
    Create a plot of the covariance confidence ellipse of x and y

    :param x: array-like, shape (n,)
    :param y: array-like, shape (n,) input data
    :param ax:  matplotlib.axes.Axes
    :param n_std: float
    The number of standard deviations to determine the ellipse's radiuses.
    :param facecolor:
    :param kwargs:
    :return:  matplotlib.patches.Ellipse
    """

    if x.size != y.size:
        raise ValueError('x and y must be the same size')

    cov = np.cov(x, y)
    pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
    # using a special case to obtain the eigenvalues of this 2D dataset
    ell_radius_x = np.sqrt(1 + pearson)
    ell_radius_y = np.sqrt(1 - pearson)
    ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2,
                      facecolor=facecolor, **kwargs)
    # Calculate the standard deviation of x from
    # the squareroot of the variance and multiplying
    # with the given number of standard deviations
    scale_x, scale_y = np.sqrt(cov[0, 0]) * n_std, np.sqrt(cov[1, 1]) * n_std
    mean_x, mean_y = np.mean(x), np.mean(y)

    transf = transforms.Affine2D().rotate_deg(45). \
        scale(scale_x, scale_y).translate(mean_x, mean_y)
    ellipse.set_transform(transf + ax.transData)
    return ax.add_patch(ellipse)


def get_correlated_dataset(n, dependency, mu, scale):
    latent = np.random.randn(n, 2)
    dependent = latent.dot(dependency)
    scaled = dependent * scale
    scaled_with_offset = scaled + mu
    # return x and y of the new, correlated dataset
    return scaled_with_offset[:, 0], scaled_with_offset[:, 1]


if __name__ == '__main__':
    fig, ax_nstd = plt.subplots(figsize=(6, 6))
    dependency_nstd = [[0.8, 0.75], [-0.2, 0.35]]
    mu = 2, -10
    scale = 8, 5
    ax_nstd.axvline(c='grey', lw=1)
    ax_nstd.axhline(c='grey', lw=1)

    # x, y = get_correlated_dataset(500, dependency_nstd, mu, scale)
    label, X = dataset.load_data(8)
    cla = 6

    x, y = X[label == cla, 0], X[label == cla, 1]
    ax_nstd.scatter(x, y)

    confidence_ellipse(x, y, ax_nstd, n_std=1,
                       label=r'$1\sigma$', edgecolor='firebrick')
    confidence_ellipse(x, y, ax_nstd, n_std=2,
                       label=r'$2\sigma$', edgecolor='fuchsia', linestyle='--')
    confidence_ellipse(x, y, ax_nstd, n_std=3,
                       label=r'$3\sigma$', edgecolor='blue', linestyle=':')
    confidence_ellipse(x, y, ax_nstd, n_std=6,
                       label=r'$5\sigma$', edgecolor='red', linestyle=':')

    ax_nstd.scatter(mu[0], mu[1], c='red', s=3)
    ax_nstd.set_title('Different standard deviations')
    ax_nstd.legend()
    plt.show()
